#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>

#define TILE_WIDTH 16
#define BLOCK_SIZE 256

// Function to generate random half precision numbers
void randomInitHalf(half *data, int size) {
    for (int i = 0; i < size; ++i) {
        data[i] = __float2half(rand() / (float)RAND_MAX);
    }
}

// Matrix multiplication kernel - GPU version with half precision
__global__ void matrixMulGPU(half *A, half *B, half *C, int width) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y * blockDim.y + threadIdx.y;

    if (col < width && row < width) {
        half sum = __float2half(0.0f);
        for (int k = 0; k < width; ++k) {
            sum += A[row * width + k] * B[k * width + col];
        }
        C[row * width + col] = sum;
    }
}

using namespace nvcuda;

__global__ void matrixMulWMMA(half *A, half *B, half *C, int width) {
    int warpCol = threadIdx.x % 16;
    int warpRow = threadIdx.x / 16;
    int warpIdx = threadIdx.y;
    int blockRow = blockIdx.y * TILE_WIDTH + warpRow;
    int blockCol = blockIdx.x * TILE_WIDTH + warpCol;
    int offsetA = blockRow * width + warpIdx * 16;
    int offsetB = warpIdx * width + blockCol;

    // Declare fragment pointers
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, half> acc_frag;
    wmma::fill_fragment(acc_frag, 0.0f);

    // Load data into fragments
    wmma::load_matrix_sync(a_frag, &A[offsetA], width);
    wmma::load_matrix_sync(b_frag, &B[offsetB], width);

    // Perform matrix multiplication
    wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);

    // Store the result back to global memory
    wmma::store_matrix_sync(&C[blockRow * width + blockCol], acc_frag, width, wmma::mem_row_major);
}

int main() {
    int width = 1024;
    int size = width * width * sizeof(half);
    half *h_A = (half*)malloc(size);
    half *h_B = (half*)malloc(size);
    half *h_C_GPU = (half*)malloc(size);
    half *h_C_WMMA = (half*)malloc(size);

    // Initialize matrices h_A and h_B with random data
    randomInitHalf(h_A, width * width);
    randomInitHalf(h_B, width * width);

    half *d_A, *d_B, *d_C_GPU;
    half *d_C_WMMA;
    cudaMalloc((void**)&d_A, size);
    cudaMalloc((void**)&d_B, size);
    cudaMalloc((void**)&d_C_GPU, size);
    cudaMalloc((void**)&d_C_WMMA, size);

    // Copy matrices from host to device
    cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice);

    dim3 blockDim(TILE_WIDTH, TILE_WIDTH);
    dim3 gridDim((width + blockDim.x - 1) / blockDim.x, (width + blockDim.y - 1) / blockDim.y);

    // Start timer for GPU version
    cudaEvent_t startGPU, stopGPU;
    cudaEventCreate(&startGPU);
    cudaEventCreate(&stopGPU);
    cudaEventRecord(startGPU);

    // Launch GPU version kernel
    matrixMulGPU<<<gridDim, blockDim>>>(d_A, d_B, d_C_GPU, width);

    // Stop timer for GPU version
    cudaEventRecord(stopGPU);
    cudaEventSynchronize(stopGPU);
    float elapsedTimeGPU;
    cudaEventElapsedTime(&elapsedTimeGPU, startGPU, stopGPU);
    printf("GPU version time: %f ms\n", elapsedTimeGPU);

    // Start timer for WMMA version
    cudaEvent_t startWMMA, stopWMMA;
    cudaEventCreate(&startWMMA);
    cudaEventCreate(&stopWMMA);
    cudaEventRecord(startWMMA);

    // Launch WMMA version kernel
    matrixMulWMMA<<<gridDim, blockDim>>>(d_A, d_B, d_C_WMMA, width);

    // Wait for kernel to finish
    cudaDeviceSynchronize();

    // Stop timer for WMMA version
    cudaEventRecord(stopWMMA);
    cudaEventSynchronize(stopWMMA);
    float elapsedTimeWMMA;
    cudaEventElapsedTime(&elapsedTimeWMMA, startWMMA, stopWMMA);
    printf("WMMA version time: %f ms\n", elapsedTimeWMMA);

    // Copy result matrices from device to host
    cudaMemcpy(h_C_GPU, d_C_GPU, size, cudaMemcpyDeviceToHost);
    cudaMemcpy(h_C_WMMA, d_C_WMMA, size, cudaMemcpyDeviceToHost);

    // Free device memory
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C_GPU);
    cudaFree(d_C_WMMA);

    // Free host memory
    free(h_A);
    free(h_B);
    free(h_C_GPU);
    free(h_C_WMMA);

    return 0;
}